package cn.jingyinghui.service;

import cn.jingyinghui.config.ArkConfig;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.schedulers.Schedulers;
import jakarta.annotation.PreDestroy;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

@Service
public class DoubaoNormalService {

    @Autowired
    private ArkService arkService;

    private String modelName;

    public Flux<String> getStreamResponse(String userMessageContent) {
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.SYSTEM)
                .content(DoubaoTuiliService.buildPrompt(userMessageContent))
                .build());
        messages.add(ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(userMessageContent)
                .build());

        var request = ChatCompletionRequest.builder()
                .model(getModelName())
                .messages(messages)
                .build();

        return Flux.from(arkService.streamChatCompletion(request)
                .doOnError(Throwable::printStackTrace)
                .subscribeOn(Schedulers.io())
                .observeOn(Schedulers.single()))
                .onErrorMap(throwable -> {
                    System.err.println("Error in chat completion: " + throwable.getMessage());
                    throwable.printStackTrace();
                    return new RuntimeException("Chat completion failed: " + throwable.getMessage());
                })
                .flatMap(response -> {
                    if (response == null || response.getChoices() == null) {
                        return Flux.error(new RuntimeException("Invalid response received"));
                    }
                    List<String> contentList = new ArrayList<>();
                    response.getChoices().forEach(choice -> {
                        if (choice.getMessage() != null && choice.getMessage().getContent() != null) {
                            contentList.add((String) choice.getMessage().getContent());
                        }
                    });
                    return contentList.isEmpty() ? 
                           Flux.error(new RuntimeException("No content received")) :
                           Flux.fromIterable(contentList);
                })
                .doOnComplete(() -> System.out.println("Stream completed successfully"))
                .doOnError(e -> System.err.println("Stream error occurred: " + e.getMessage()));
    }

    @PreDestroy
    public void onDestroy() {
        try {
            System.out.println("Shutting down ArkService...");
            CompletableFuture.runAsync(() -> {
                arkService.shutdownExecutor();
            }).get(5, TimeUnit.SECONDS);
            System.out.println("ArkService shutdown completed.");
        } catch (Exception e) {
            System.err.println("Error shutting down ArkService: " + e.getMessage());
        }
    }

    public String getModelName() {
        return modelName;
    }

    @Value("${yiyi.daoubao.normal.model.name}")
    public void setModelName(String modelName) {
        this.modelName = modelName;
    }
}
